import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data

mnist = input_data.read_data_sets('MNIST_data',one_hot=True)
batch_size =100
n_batch = mnist.train.num_examples // batch_size

# 初始化权值 截断的正态分布
def weight_variable(shape):
    initial = tf.truncated_normal(shape,stddev=0.1)
    return tf.Variable(initial)

# 初始化偏置
def biase_variable(shape):
    initial =tf.constant(0.1,shape=shape)
    return  tf.Variable(initial)

# 卷积层
# input：指需要做卷积的输入图像，[训练时一个batch的图片数量, 图片高度, 图片宽度, 图像通道数(黑白是1，彩色3)]，注意这是一个4维的Tensor，要求类型为float32和float64其中之一
# filter：相当于CNN中的卷积核，[卷积核的高度，卷积核的宽度，图像通道数，卷积核个数]
# strides：卷积时在图像每一维的步长，这是一个一维的向量，长度4，0和3 是1，1是x，2是y
# padding：string类型的量，只能是"SAME","VALID"
def conv2d(x,W):
    # x
    return tf.nn.conv2d(x,W,strides=[1,1,1,1],padding='SAME')

# 池化层
def max_pool_2x2(x):
    return tf.nn.max_pool(x,ksize=[1,2,2,1],strides=[1,2,2,1],padding='SAME')

# 定义PLACEHOLDER
x = tf.placeholder(tf.float32,[None,784])
y = tf.placeholder(tf.float32,[None,10])

# 改变x的格式转为4D的向量[训练时一个batch的图片数量, 图片高度, 图片宽度, 图像通道数}
x_image = tf.reshape(x,[-1,28,28,1])

# 初始化第一个卷积层的W和B
# 5*5采样窗口，32个卷积核从一个平面抽取，输出32
W_conv1 = weight_variable([5,5,1,32])
B_conv1 = biase_variable([32])

# x_image的权值向量进行卷积 加上偏置  应用于激活函数
h_conv1 = tf.nn.relu(conv2d(x_image,W_conv1)+B_conv1)
h_pool1 = max_pool_2x2(h_conv1)

# 初始化第2个卷积层的W和B
# 5*5采样窗口，32个特征 输出是64个卷积核从32平面抽取
W_conv2 = weight_variable([5,5,32,64])
B_conv2 = biase_variable([64])

# x_image的权值向量进行卷积 加上偏置  应用于激活函数
h_conv2 = tf.nn.relu(conv2d(h_pool1,W_conv2)+B_conv2)
h_pool2 = max_pool_2x2(h_conv2)

#28*28 第一次卷积后28*28（same），池化后14*14
# 第二次卷积后 14*14，池化 7*7
# 输出是 64个 7*7的特征平面[

# 初始化全连接层的权值
# 上一层7*7*64个神经元，全连接层1024个神经元  1024个结点
W_fc1 = weight_variable([7*7*64,1024])
b_fc1 = biase_variable([1024])

# 池化层2 输出扁平化1维
h_pool2_flat = tf.reshape(h_pool2,[-1,7*7*64])
h_fc1 = tf.nn.relu(tf.matmul(h_pool2_flat,W_fc1)+b_fc1)

# 输出概率
keep_prob = tf.placeholder(tf.float32)
h_fc1_drop = tf.nn.dropout(h_fc1,keep_prob)

# 初始化全连接层2
W_fc2 = weight_variable([1024,10])
b_fc2 = biase_variable([10])

# 计算输出
prediction = (tf.matmul(h_fc1_drop,W_fc2)+b_fc2)

# 交叉熵
cross_entropy = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(labels=y,logits=prediction))

# Adam
train_step = tf.train.AdamOptimizer(1e-4).minimize(cross_entropy)

# 结果放在布尔列表
correct_prediction = tf.equal(tf.argmax(prediction,1),tf.argmax(y,1))

accuracy = tf.reduce_mean(tf.cast(correct_prediction,tf.float32))

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    for epoch in range(21):
        for batch in range(n_batch):
            batch_xs,batch_ys = mnist.train.next_batch(batch_size)
            sess.run(train_step,feed_dict={x:batch_xs,y:batch_ys,keep_prob:0.7})

        acc = sess.run(accuracy,feed_dict={x: mnist.train.images, y: mnist.train.labels,keep_prob:1.0})
        print('Iter' + str(epoch) + ',Testing Accuracy:' + str(acc) )
